/*
 *  NavierStokesSolver.cpp
 *  NSESolver
 *
 *  Created by Patrick Zulian on 3/3/11.
 *  Copyright 2011 ICS - USI - Lugano. All rights reserved.
 *
 */

#include <limits>

#include "Utils.h"

#include "NavierStokesSolver.h"


#include "InitialValue.h"
#include "BoundaryHandler.h"

#define NSE_NUM_BOUND 10e10

//#define NSE_ASSERT

//Dirty Macros to check bound of the data
#ifdef NSE_ASSERT

#define NSE_ASSERT_VALID(val) \
if(val >=  std::numeric_limits< double >::max() || val <= -std::numeric_limits< double >::max()) { \
std::cout << "invalid = " << val << std::endl; \
NSE_SOLVER_ASSERT( val <  std::numeric_limits< double >::max() ); \
NSE_SOLVER_ASSERT( val > -std::numeric_limits< double >::max() ); } \
else if(val > NSE_NUM_BOUND || val < -NSE_NUM_BOUND )  {\
std::cout << "too big = " << val << std::endl; \
NSE_SOLVER_ASSERT( val <= NSE_NUM_BOUND ); \
NSE_SOLVER_ASSERT( val >= -NSE_NUM_BOUND ); }

#else
#define NSE_ASSERT_VALID(val) (void)val

#endif //NSE_ASSERT


namespace nsesolver {
    
    //to wrap the operator for gdb
	double GetValAt(const Matrixd &u, const int i, const int j);
    double GetValAt(const Matrixd &u, const int i, const int j) {
        return u(i,j);
    }
    
    NavierStokesSolver::NavierStokesSolver()
    : writer()
    { }
	
    
	NavierStokesSolver::~NavierStokesSolver()
	{ }
	
	void NavierStokesSolver::solve(CFDProblem &problem,const double t0,const double tEnd,const double omega,const double Re,const double gx,const double gy)
	{
        
		using std::min;
		using std::max;
        
        using std::cout;
        using std::endl;
        
        
        using std::vector;
		
        
        
        if (!writer.isNull()) {
            writer->solverStarted(problem);
        }
        
        
        
        double t = t0;
        
		
        const double tauMax = 0.5;//0.015;
        double tau = tauMax;
		
        const double dx=problem.DX();
        const double dy=problem.DY();
        
		const double dxSquared = dx*dx;
		const double dySquared = dy*dy;
        
        const double ReInv=1.0/Re;
        const double ReScaled = Re/2.0 *  1.0/ (1.0/dxSquared + 1.0/dySquared);
		
		Matrixd & u = problem.U();
		Matrixd & v = problem.V();
		//Matrixd & p = problem.P();
        
		Matrixd F(u.xmax(), u.ymax());
		Matrixd G(v.xmax(), v.ymax());
		
		Matrixd rhs(problem.IMax()+1, problem.JMax()+1);
        bool isRoot = problem.getTopology().getCommunicator().isRoot();
        
        //!!!IVP has to satisfy eq (2c). Continuity equation
        //TODO
        problem.setInitialValues();
        
        F.allSet(0);
        G.allSet(0);
        rhs.allSet(0);
        
        
        
        
        int q = 1;
        vector< double > ts;
        vector< double > taus;
        ts.push_back(t);
        taus.push_back(tau);
        
        //Set boundary conditions for u^n and v^n according to section 3.2
        problem.setBoundary();
        
		while (t < tEnd) {
            
			//Compute tau 
			tau = min( min( min(ReScaled, dx/problem.uAbsMax() ), dy/problem.vAbsMax() ), tauMax);            
            
            //Compute F^n, G^n by means of (26) and (27) and Set the boundaries on F, G by means of (35)
			computeFandG(F, G, problem, tau, ReInv, gx, gy );
            
            //Compute rhs for pressure equation by means of the right hans side of equation (28)
            computeRHS(rhs, problem, F, G, tau);
            
            
            //Set p^n as the initial value for the SOR scheme
            //Done internally in SORStep(...)
            SORSolve(problem, rhs, omega);
            
            //compute u_n+1 v_n+1 by (24) (25)
            updateVelocities(problem, G, F, tau);
            
            problem.setBoundary();
            
            //std::cout<<"hereu:\n"<<u(Range(2, 6),Range(2, 7))<<std::endl;
            //std::cout<<"herev:\n"<<v(Range(2, 7),Range(2, 6))<<std::endl;
            
            t += tau;
            if(isRoot)
                cout << "t = " << t << endl;
            if(!writer.isNull())
                writer->solverStep(problem, q++);
            
            //exit(0);
            ts.push_back(t);
            taus.push_back(tau);
            
		}
        
        if(!writer.isNull()) {
            writer->solverEnded(ts, taus);
        }
        
	}
    
    
    /// Compute F^n, G^n by means of (26) and (27) and Set the boundaries on F, G by means of (35)
    void  NavierStokesSolver::computeFandG(Matrixd &F, Matrixd &G, CFDProblem & problem, const double tau, const double ReInv, const double gx, const double gy)
    {
        const Matrixd &u = problem.U();
        const Matrixd &v = problem.V();
        
        //Compute F^n, G^n by means of (26) and (27)
        for(int i=1;i<=problem.IMax()-1;i++) {
            for(int j=1;j<=problem.JMax();++j) {
                F(i,j)=u(i,j)+tau*(ReInv*
                                   ( problem.d2udx2(i, j)+ problem.d2udy2(i, j) ) -
                                   problem.duSquaredx(i, j) -
                                   problem.duvdy(i, j) + gx);
                
                NSE_ASSERT_VALID( F(i, j) );
            }
        }
        
        for(int i=1;i<=problem.IMax();i++) {
            for(int j=1;j<=problem.JMax()-1;++j) {
                G(i,j)=v(i,j)+tau*(ReInv*
                                   ( problem.d2vdx2(i,j)+problem.d2vdy2(i,j) ) - 
                                   problem.dvSquaredy(i,j) - 
                                   problem.duvdx(i, j) + gy);
                
                NSE_ASSERT_VALID( G(i, j) );
            }
        }
        
        problem.setFGBoundary(F, G);
        
    }
    
    ///compute u_n+1 v_n+1 by (24) (25)
    void NavierStokesSolver::updateVelocities(CFDProblem &problem, const Matrixd &G, const Matrixd &F, const double tau)
    {
        
        Matrixd &u = problem.U();
        Matrixd &v = problem.V();
        
        for(int i=1;i<=problem.IMax()-1;++i){
            for(int j=1;j<=problem.JMax();++j) {
                NSE_ASSERT_VALID( u(i,j) );
                NSE_ASSERT_VALID( F(i,j) );
                u(i,j)= F(i,j) - tau * problem.dpdx(i,j);
                NSE_ASSERT_VALID( u(i,j) );
            }
        }
        
        for(int i=1;i<=problem.IMax();++i){
            for(int j=1;j<=problem.JMax()-1;++j) {
                NSE_ASSERT_VALID( v(i,j) );
                v(i,j) = G(i,j) - tau * problem.dpdy(i,j);
                NSE_ASSERT_VALID( G(i,j) );
                NSE_ASSERT_VALID( v(i,j) );
            }
        }
        
    }
    
    
    /// Compute rhs for pressure equation by means of the right hans side of equation (28)
    void  NavierStokesSolver::computeRHS(Matrixd &rhs, const CFDProblem &problem, const Matrixd F, const Matrixd G, const double tau )
    {
        for(int i=1;i<=problem.IMax();++i) {
            for(int j=1;j<=problem.JMax();++j) {
                rhs(i,j)= 1.0/tau * ( ( F(i,j) - F(i-1,j) )/problem.DX() + 
                                     (  G(i,j) - G(i,j-1) )/problem.DY() );
                
                NSE_ASSERT_VALID( rhs(i,j) );
            }
        }
        
    }
    
    
    
    void NavierStokesSolver::SORSolve(CFDProblem &problem, const Matrixd &rhs, const double omega)
    {
        //Compute residual by means of (29)
        double res = computeResidual(problem, rhs);
        double prevRes = res;
        
        //Set p^n as the initial value for the SOR scheme
        //Done internally in SORStep(...)
        
        //Compute boundary values by means of (34)
        problem.setPBoundary();
        
        while ( res > Numerics::eps< float >() ) {
            //Compute an SOR step emplyoing Algorithm 2
            SORStep(problem, rhs, omega);
            
            
            //Compute boundary values by means of (34)
            problem.setPBoundary();
            
            //Compute the residual by means of (29)
            res = computeResidual(problem, rhs);
            if (res > prevRes && problem.getTopology().getCommunicator().isRoot()) {
                std::cout << "res: " << res << std::endl;
            }
            
            prevRes = res;
        }
    }
    
    
    void NavierStokesSolver::SORStep(CFDProblem &problem, const Matrixd &rhs, const double omega)
    {
        Matrixd &p(problem.P());
        
        const double inv2dx=1.0/Math::Pow(problem.DX(),(int)2);
        const double inv2dy=1.0/Math::Pow(problem.DY(),(int)2);
        
        NSE_ASSERT_VALID( inv2dx );
        NSE_ASSERT_VALID( inv2dy ); 
        
        const double dumpDiv = omega/(2.0 * (inv2dx + inv2dy) );
        NSE_ASSERT_VALID( dumpDiv );
        
        
        for (int j=1; j <= problem.JMax(); ++j) {
            for (int i=1; i <= problem.IMax(); ++i) {
                if (problem.getBoundaryConditions().getObstacle().isObstacleCell(i, j)) {
                    continue;
                }
                
                NSE_ASSERT_VALID( p(i, j)  );
                NSE_ASSERT_VALID( p(i+1, j)  );
                NSE_ASSERT_VALID( p(i-1, j)  );
                NSE_ASSERT_VALID( p(i, j+1)  );
                NSE_ASSERT_VALID( p(i, j-1)  );
                
                NSE_ASSERT_VALID( rhs(i, j)   );
                
                double val = (1.0 - omega) * p(i,j) + dumpDiv*(
                                                               ( p(i+1,j) + p(i-1,j) ) * inv2dx +
                                                               ( p(i,j+1) + p(i,j-1) ) * inv2dy -
                                                               rhs(i,j) 
                                                               );
                
                NSE_ASSERT_VALID( val );
                p(i, j) = val;
                
            }
        }
    }
	
	double NavierStokesSolver::computeResidual(const CFDProblem &problem, const Matrixd &rhs) const
	{
        const Obstacle & obs(problem.getBoundaryConditions().getObstacle());
        
		double res = 0;
        
        //DEBUG
//        static int q = 0;
//        Matrixd residual(problem.P().xmax()-1, problem.P().ymax()-1);
//        residual.allSet(0);
		
        for(int j=1;j<=problem.JMax();++j) {
            for(int i=1;i<=problem.IMax();++i) {
                if (obs.isObstacleCell(i, j)) {
                    continue;
                }
                
                const double resij = problem.d2pdx2(i, j) + problem.d2pdy2(i,j) - rhs(i,j);
				res += Math::Pow(resij, static_cast< int >(2));
                
                //DEBUG
//                residual(i-1, j-1) = fabs(resij);
            }
        }
		
		res /= problem.IMax() * problem.JMax();
		
        problem.getTopology().getCommunicator().allReduce(&res, 1, MPI_SUM);
		res = Math::Pow(res, 0.5);
        
        //DEBUG
//        if (!writer.isNull()) {
//            std::stringstream ss;
//            ss << "residual" << q++;
//            writer->writeMatrix(problem, residual, "residual", ss.str());
//        }
        
		return res;
	}
    
    
	
	
}
